#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
洪水预报参数验证工具

基于LangChain工具框架实现的参数验证功能
"""

from langchain.tools import BaseTool
from typing import Dict, Any, Optional, List, Union
from pydantic import BaseModel, Field, validator
import re
import json
from datetime import datetime, timedelta
from loguru import logger


class ParameterValidatorInput(BaseModel):
    """参数验证工具输入模型"""
    parameters: Union[Dict[str, Any], str] = Field(description="待验证的参数字典或JSON字符串")
    validation_rules: Optional[Dict[str, Any]] = Field(default=None, description="验证规则")
    
    @validator('parameters', pre=True)
    def parse_parameters(cls, v):
        """预处理参数，将字符串转换为字典"""
        if isinstance(v, str):
            try:
                parsed = json.loads(v)
                if isinstance(parsed, dict):
                    return parsed
                else:
                    raise ValueError("解析的JSON不是字典格式")
            except json.JSONDecodeError:
                raise ValueError(f"无法解析JSON字符串: {v}")
        elif isinstance(v, dict):
            return v
        else:
            raise ValueError(f"不支持的输入类型: {type(v)}")
    
    @classmethod
    def parse_input(cls, input_data):
        """智能解析输入数据"""
        if isinstance(input_data, str):
            try:
                # 尝试解析JSON字符串
                parsed_data = json.loads(input_data)
                # 确保解析出的是字典类型
                if not isinstance(parsed_data, dict):
                    raise ValueError("解析的JSON不是字典格式")
                # 如果解析出的是字典，检查是否包含parameters字段
                if 'parameters' in parsed_data:
                    # 提取parameters字段的值作为实际参数
                    return cls(parameters=parsed_data['parameters'])
                else:
                    # 如果没有parameters字段，假设整个字典就是参数
                    return cls(parameters=parsed_data)
            except json.JSONDecodeError:
                raise ValueError(f"无法解析JSON字符串: {input_data}")
        elif isinstance(input_data, dict):
            # 如果输入是字典，检查是否包含parameters字段
            if 'parameters' in input_data:
                # 提取parameters字段的值作为实际参数
                return cls(parameters=input_data['parameters'])
            else:
                # 如果没有parameters字段，假设整个字典就是参数
                return cls(parameters=input_data)
        else:
            raise ValueError(f"不支持的输入类型: {type(input_data)}")

class ParameterValidatorTool(BaseTool):
    """洪水预报参数验证工具
    
    验证提取的洪水预报参数的格式和合理性
    """
    
    name: str = "ParameterValidatorTool"
    description: str = "验证提取的洪水预报参数的格式和合理性，包括数值范围、时间格式、必填字段等"
    args_schema: type = ParameterValidatorInput
    
    def __init__(self):
        super().__init__()
        self._init_validation_rules()
    
    def _init_validation_rules(self):
        """初始化验证规则"""
        self._validation_rules = {
            # 数值范围验证
            'numerical_ranges': {
                'forecast_duration': {'min': 1, 'max': 168, 'unit': '小时'},  # 1小时到7天
                'rainfall_forecast_period': {'min': 1, 'max': 72, 'unit': '小时'},  # 1小时到3天
                'water_level': {'min': 0, 'max': 1000, 'unit': '米'},
                'flow_rate': {'min': 0, 'max': 100000, 'unit': '立方米每秒'},
                'rainfall': {'min': 0, 'max': 1000, 'unit': '毫米'}
            },
            
            # 必填字段
            'required_fields': ['reservoir_name', 'forecast_model', 'start_time', 
                              'forecast_duration', 'rainfall_forecast_period', 'target_variable'],
            
            # 水库白名单（可配置）
            'reservoir_whitelist': [
                '大伙房水库', '三峡水库', '葛洲坝水库', '丹江口水库', '小浪底水库',
                '龙羊峡水库', '刘家峡水库', '万家寨水库', '天生桥水库', '岩滩水库',
                '龙滩水库', '百色水库', '紫坪铺水库', '瀑布沟水库', '糯扎渡水库',
                '小湾水库', '景洪水库', '漫湾水库', '大朝山水库', '功果桥水库'
            ],
            
            # 预报模型白名单
            'model_whitelist': [
                '新安江', '新安江模型', '萨克拉门托', '萨克拉门托模型',
                'tank', 'Tank', 'TANK', 'HBV', 'SWAT', 'VIC', 'TOPMODEL',
                '神经网络', '机器学习', '深度学习', 'LSTM', 'CNN', 'RNN'
            ],
            
            # 目标变量白名单
            'target_variable_whitelist': [
                '流量', '水位', '库容', '入库流量', '出库流量', '蓄水量'
            ],
            
            # 时间格式模式
            'time_format_patterns': [
                r'^\d{4}-\d{1,2}-\d{1,2}\s+\d{1,2}:\d{1,2}$',  # YYYY-MM-DD HH:MM
                r'^\d{4}-\d{1,2}-\d{1,2}$',  # YYYY-MM-DD
            ]
        }
    
    def _run(self, parameters: Dict[str, Any], validation_rules: Optional[Dict[str, Any]] = None) -> str:
        """执行参数验证"""
        try:
            # 第一步：安全解析输入参数
            parse_result = self._safe_parse_parameters(parameters)
            if not parse_result['success']:
                return json.dumps({
                    'all_valid': False,
                    'validation_errors': {'input_format': parse_result['error']},
                    'validation_warnings': {},
                    'validated_parameters': {},
                    'validation_score': 0.0,
                    'missing_parameters': [],
                    'invalid_parameters': ['input_format'],
                    'user_guidance': parse_result['guidance']
                }, ensure_ascii=False, indent=2)
            
            actual_parameters = parse_result['data']
            logger.info(f"开始验证参数: {actual_parameters}")
            
            # 使用自定义规则或默认规则
            rules = validation_rules or self._validation_rules
            
            validation_result = {
                'all_valid': True,
                'validation_errors': {},
                'validation_warnings': {},
                'validated_parameters': actual_parameters.copy(),
                'validation_score': 0.0,
                'missing_parameters': [],
                'invalid_parameters': [],
                'user_guidance': []  # 新增：用户指导信息
            }
            
            # 执行各项验证
            self._validate_required_fields(actual_parameters, rules, validation_result)
            self._validate_reservoir_name(actual_parameters, rules, validation_result)
            self._validate_forecast_model(actual_parameters, rules, validation_result)
            self._validate_start_time(actual_parameters, rules, validation_result)
            self._validate_numerical_ranges(actual_parameters, rules, validation_result)
            self._validate_target_variable(actual_parameters, rules, validation_result)
            self._validate_logical_consistency(actual_parameters, validation_result)
            
            # 计算验证分数
            validation_result['validation_score'] = self._calculate_validation_score(validation_result)
            
            # 生成验证摘要
            validation_result['validation_summary'] = self._generate_validation_summary(validation_result)
            
            logger.info(f"参数验证完成: {validation_result}")
            return json.dumps(validation_result, ensure_ascii=False, indent=2)
            
        except Exception as e:
            logger.error(f"参数验证过程中发生错误: {str(e)}")
            error_result = {
                'all_valid': False,
                'validation_errors': {'system_error': f"验证过程中发生系统错误: {str(e)}"},
                'validation_warnings': {},
                'validated_parameters': parameters if isinstance(parameters, dict) else {},
                'validation_score': 0.0,
                'missing_parameters': [],
                'invalid_parameters': [],
                'user_guidance': [
                    "系统在验证参数时遇到了内部错误。请检查参数格式是否正确，或联系系统管理员。",
                    "建议：确保所有参数都是有效的字典格式或JSON字符串。"
                ]
            }
            return json.dumps(error_result, ensure_ascii=False, indent=2)
    
    def _validate_required_fields(self, params: Dict[str, Any], rules: Dict[str, Any], result: Dict[str, Any]):
        """验证必填字段"""
        required_fields = rules.get('required_fields', [])
        
        # 定义字段说明和示例
        field_descriptions = {
            'reservoir_name': {'desc': '水库名称', 'example': '大伙房水库'},
            'forecast_model': {'desc': '预报模型', 'example': '新安江模型'},
            'start_time': {'desc': '预报开始时间', 'example': '2024-01-15 08:00'},
            'forecast_duration': {'desc': '预报时长（小时）', 'example': '48'},
            'rainfall_forecast_period': {'desc': '降雨预见期（小时）', 'example': '12'},
            'target_variable': {'desc': '预报目标变量', 'example': '流量'}
        }
        
        for field in required_fields:
            if field not in params or params[field] is None or params[field] == '':
                result['missing_parameters'].append(field)
                field_info = field_descriptions.get(field, {'desc': field, 'example': '请提供有效值'})
                result['validation_errors'][field] = f"缺少必填字段: {field_info['desc']}"
                result['user_guidance'].append(
                    f"请提供{field_info['desc']}（{field}），例如：\"{field}\": \"{field_info['example']}\""
                )
                result['all_valid'] = False
    
    def _validate_reservoir_name(self, params: Dict[str, Any], rules: Dict[str, Any], result: Dict[str, Any]):
        """验证水库名称"""
        reservoir_name = params.get('reservoir_name')
        if not reservoir_name:
            return
        
        whitelist = rules.get('reservoir_whitelist', [])
        if whitelist and reservoir_name not in whitelist:
            result['validation_warnings']['reservoir_name'] = f"水库名称 '{reservoir_name}' 不在已知列表中，请确认是否正确"
            result['user_guidance'].append(
                f"水库名称 '{reservoir_name}' 不在系统已知列表中。如果确认无误请继续，否则请从以下列表中选择：{', '.join(whitelist[:5])}等"
            )
        
        # 检查名称格式
        if not re.match(r'^[\u4e00-\u9fa5]+(?:水库|库区|坝|水电站)?$', reservoir_name):
            result['validation_errors']['reservoir_name'] = f"水库名称格式不正确: {reservoir_name}"
            result['invalid_parameters'].append('reservoir_name')
            result['user_guidance'].append(
                f"水库名称格式不正确。请使用中文名称，可以包含'水库'、'库区'、'坝'、'水电站'等后缀。正确示例：'大伙房水库'、'三峡水库'"
            )
            result['all_valid'] = False
    
    def _validate_forecast_model(self, params: Dict[str, Any], rules: Dict[str, Any], result: Dict[str, Any]):
        """验证预报模型"""
        forecast_model = params.get('forecast_model')
        if not forecast_model:
            return
        
        whitelist = rules.get('model_whitelist', [])
        if whitelist:
            # 直接匹配
            if forecast_model in whitelist:
                return
            
            # 智能匹配：去除"模型"后缀后再匹配
            model_name_cleaned = forecast_model
            if forecast_model.endswith('模型'):
                model_name_cleaned = forecast_model[:-2]  # 去除"模型"两个字
            
            # 检查去除"模型"后的名称是否在白名单中
            if model_name_cleaned in whitelist:
                return
            
            # 反向匹配：检查白名单中是否有项目以当前模型名开头
            for white_model in whitelist:
                if white_model.startswith(model_name_cleaned) or model_name_cleaned.startswith(white_model):
                    return
            
            # 如果都不匹配，给出警告
            result['validation_warnings']['forecast_model'] = f"预报模型 '{forecast_model}' 不在已知列表中，请确认是否正确"
            result['user_guidance'].append(
                f"预报模型 '{forecast_model}' 不在系统已知列表中。如果确认无误请继续，否则请从以下模型中选择：{', '.join(whitelist[:5])}等"
            )
    
    def _validate_start_time(self, params: Dict[str, Any], rules: Dict[str, Any], result: Dict[str, Any]):
        """验证开始时间"""
        start_time = params.get('start_time')
        if not start_time:
            return
        
        # 检查时间格式
        time_patterns = rules.get('time_format_patterns', [])
        format_valid = False
        
        for pattern in time_patterns:
            if re.match(pattern, start_time):
                format_valid = True
                break
        
        if not format_valid:
            result['validation_errors']['start_time'] = f"时间格式不正确: {start_time}"
            result['invalid_parameters'].append('start_time')
            result['user_guidance'].append(
                f"时间格式不正确。请使用以下格式之一：\n1. 'YYYY-MM-DD HH:MM'（如：'2024-01-15 08:00'）\n2. 'YYYY-MM-DD'（如：'2024-01-15'）"
            )
            result['all_valid'] = False
            return
        
        # 检查时间合理性
        try:
            if ' ' in start_time:
                parsed_time = datetime.strptime(start_time, '%Y-%m-%d %H:%M')
            else:
                parsed_time = datetime.strptime(start_time, '%Y-%m-%d')
            
            now = datetime.now()
            
            # 检查是否为未来时间（允许过去24小时内的时间）
            if parsed_time < now - timedelta(hours=24):
                result['validation_warnings']['start_time'] = f"开始时间 '{start_time}' 距离现在较远，请确认是否正确"
                result['user_guidance'].append(
                    f"开始时间 '{start_time}' 距离当前时间较远。通常预报开始时间应该在当前时间前后24小时内。请确认时间是否正确。"
                )
            
            # 检查是否过于遥远的未来
            if parsed_time > now + timedelta(days=30):
                result['validation_warnings']['start_time'] = f"开始时间 '{start_time}' 距离现在过远，请确认是否正确"
                result['user_guidance'].append(
                    f"开始时间 '{start_time}' 距离当前时间过远（超过30天）。请确认是否需要如此长期的预报。"
                )
                
        except ValueError:
            result['validation_errors']['start_time'] = f"时间格式解析失败: {start_time}"
            result['invalid_parameters'].append('start_time')
            result['user_guidance'].append(
                f"时间 '{start_time}' 无法解析。请检查日期是否有效（如月份1-12，日期1-31），时间是否有效（如小时0-23，分钟0-59）。"
            )
            result['all_valid'] = False
    
    def _validate_numerical_ranges(self, params: Dict[str, Any], rules: Dict[str, Any], result: Dict[str, Any]):
        """验证数值范围"""
        numerical_ranges = rules.get('numerical_ranges', {})
        
        for param_name, range_rule in numerical_ranges.items():
            value = params.get(param_name)
            if value is None:
                continue
            
            try:
                numeric_value = float(value)
                
                if numeric_value < range_rule['min'] or numeric_value > range_rule['max']:
                    result['validation_errors'][param_name] = (
                        f"{param_name} 值 {numeric_value} 超出合理范围 "
                        f"[{range_rule['min']}, {range_rule['max']}] {range_rule['unit']}"
                    )
                    result['invalid_parameters'].append(param_name)
                    result['user_guidance'].append(
                        f"{param_name} 的值 {numeric_value} 超出了合理范围。请提供 {range_rule['min']} 到 {range_rule['max']} {range_rule['unit']} 之间的值。"
                    )
                    result['all_valid'] = False
                    
            except (ValueError, TypeError):
                result['validation_errors'][param_name] = f"{param_name} 不是有效的数值: {value}"
                result['invalid_parameters'].append(param_name)
                result['user_guidance'].append(
                    f"{param_name} 的值 '{value}' 不是有效的数值。请提供数字格式的值，例如：48、12.5 等。"
                )
                result['all_valid'] = False
    
    def _validate_target_variable(self, params: Dict[str, Any], rules: Dict[str, Any], result: Dict[str, Any]):
        """验证目标变量"""
        target_variable = params.get('target_variable')
        if not target_variable:
            return
        
        whitelist = rules.get('target_variable_whitelist', [])
        if whitelist and target_variable not in whitelist:
            result['validation_warnings']['target_variable'] = f"目标变量 '{target_variable}' 不在已知列表中，请确认是否正确"
            result['user_guidance'].append(
                f"目标变量 '{target_variable}' 不在系统已知列表中。如果确认无误请继续，否则请从以下变量中选择：{', '.join(whitelist)}等"
            )
    
    def _validate_logical_consistency(self, params: Dict[str, Any], result: Dict[str, Any]):
        """验证逻辑一致性"""
        # 检查预报时长和降雨预见期的关系
        forecast_duration = params.get('forecast_duration')
        rainfall_period = params.get('rainfall_forecast_period')
        
        if forecast_duration and rainfall_period:
            try:
                duration_val = float(forecast_duration)
                period_val = float(rainfall_period)
                
                if period_val > duration_val:
                    result['validation_warnings']['logical_consistency'] = (
                        f"降雨预见期({period_val}小时)大于预报时长({duration_val}小时)，请确认是否合理"
                    )
                    result['user_guidance'].append(
                        f"逻辑不一致：降雨预见期({period_val}小时)不应大于预报时长({duration_val}小时)。请调整参数使降雨预见期小于或等于预报时长。"
                    )
                    
            except (ValueError, TypeError):
                pass  # 数值验证在其他地方处理
    
    def _calculate_validation_score(self, result: Dict[str, Any]) -> float:
        """计算验证分数"""
        if not result['all_valid']:
            return 0.0
        
        # 基础分数
        score = 1.0
        
        # 警告扣分（每个警告扣0.1分）
        warning_count = len(result['validation_warnings'])
        warning_penalty = warning_count * 0.1
        score -= warning_penalty
        
        return max(score, 0.0)
    
    def _safe_parse_parameters(self, parameters) -> Dict[str, Any]:
        """安全解析输入参数，提供详细的错误指导"""
        if parameters is None:
            return {
                'success': False,
                'error': '输入参数为空',
                'guidance': '请提供有效的参数。参数应该是包含洪水预报信息的字典或JSON字符串。'
            }
        
        if isinstance(parameters, dict):
            # 检查字典是否为空
            if not parameters:
                return {
                    'success': False,
                    'error': '参数字典为空',
                    'guidance': '请提供包含洪水预报参数的字典，例如：{"reservoir_name": "大伙房水库", "forecast_model": "新安江模型"}'
                }
            
            # 检查是否包含嵌套的parameters字段
            if 'parameters' in parameters and isinstance(parameters['parameters'], dict):
                # 如果包含parameters字段，提取其值作为实际参数
                actual_params = parameters['parameters']
                if not actual_params:
                    return {
                        'success': False,
                        'error': '嵌套的parameters字段为空',
                        'guidance': '请在parameters字段中提供洪水预报参数，例如：{"parameters": {"reservoir_name": "大伙房水库"}}'
                    }
                return {'success': True, 'data': actual_params}
            
            return {'success': True, 'data': parameters}
        
        if isinstance(parameters, str):
            # 检查字符串是否为空
            if not parameters.strip():
                return {
                    'success': False,
                    'error': '输入字符串为空',
                    'guidance': '请提供有效的JSON字符串，例如：{"reservoir_name": "大伙房水库", "forecast_model": "新安江模型"}'
                }
            
            # 尝试解析JSON
            try:
                parsed_data = json.loads(parameters)
                if not isinstance(parsed_data, dict):
                    return {
                        'success': False,
                        'error': f'JSON解析结果不是字典格式，而是{type(parsed_data).__name__}',
                        'guidance': '请确保JSON字符串表示一个对象（字典），例如：{"reservoir_name": "大伙房水库"}，而不是数组或其他类型。'
                    }
                
                if not parsed_data:
                    return {
                        'success': False,
                        'error': '解析出的字典为空',
                        'guidance': '请在JSON对象中包含洪水预报参数，例如：{"reservoir_name": "大伙房水库", "forecast_model": "新安江模型"}'
                    }
                
                return {'success': True, 'data': parsed_data}
                
            except json.JSONDecodeError as e:
                return {
                    'success': False,
                    'error': f'JSON格式错误：{str(e)}',
                    'guidance': f'请检查JSON格式是否正确。常见问题：1) 缺少引号；2) 多余的逗号；3) 括号不匹配。正确格式示例：{{"reservoir_name": "大伙房水库"}}'
                }
        
        # 其他类型
        return {
            'success': False,
            'error': f'不支持的参数类型：{type(parameters).__name__}',
            'guidance': '参数必须是字典或JSON字符串格式。请提供包含洪水预报信息的字典或有效的JSON字符串。'
        }
    
    def _generate_validation_summary(self, result: Dict[str, Any]) -> str:
        """生成验证摘要"""
        if result['all_valid']:
            if result['validation_warnings']:
                return f"验证通过，但有 {len(result['validation_warnings'])} 个警告需要注意"
            else:
                return "所有参数验证通过"
        else:
            error_count = len(result['validation_errors'])
            missing_count = len(result['missing_parameters'])
            guidance_count = len(result.get('user_guidance', []))
            summary = f"验证失败：{error_count} 个错误，{missing_count} 个缺失参数"
            if guidance_count > 0:
                summary += f"，已提供 {guidance_count} 条修改建议"
            return summary

# 使用示例
if __name__ == "__main__":
    tool = ParameterValidatorTool()
    
    # 测试用例1：正确的参数
    print("=== 测试用例1：正确参数 ===")
    test_params_valid = {
        'reservoir_name': '大伙房水库',
        'forecast_model': '新安江模型',
        'start_time': '2024-01-15 08:00',
        'forecast_duration': 48,
        'rainfall_forecast_period': 12,
        'target_variable': '流量'
    }
    result1 = tool._run(test_params_valid)
    print(result1)
    
    # 测试用例2：有错误的参数
    print("\n=== 测试用例2：错误参数 ===")
    test_params_invalid = {
        'reservoir_name': 'InvalidReservoir123',  # 格式错误
        'start_time': '2024-13-45',  # 日期无效
        'forecast_duration': 'not_a_number',  # 非数值
    }
    result2 = tool._run(test_params_invalid)
    print(result2)
    
    # 测试用例3：JSON字符串输入
    print("\n=== 测试用例3：JSON字符串输入 ===")
    json_input = '{"reservoir_name": "三峡水库"}'
    result3 = tool._run(json_input)
    print(result3)
    
    # 测试用例4：错误的JSON格式
    print("\n=== 测试用例4：错误JSON格式 ===")
    invalid_json = '{reservoir_name: 三峡水库}'
    result4 = tool._run(invalid_json)
    print(result4)